from sympy import false

from ultralytics import YOLO
# from ultralytics.nn.modules import MobileViTv2Block
# print('MobileViTv2Block' in globals())
import torch
torch.cuda.empty_cache()
print(torch.__version__)
print(torch.version.cuda)
print(torch.backends.cudnn.version())
# 检查GPU设备
print(f"CUDA可用: {torch.cuda.is_available()}")
print(f"GPU数量: {torch.cuda.device_count()}")
for i in range(torch.cuda.device_count()):
    print(f"GPU {i}: {torch.cuda.get_device_name(i)}")

def main():
    # 加载预训练的 YOLO11n 模型
    config_path = "yolo11.yaml"


    # 验证配置文件内容
    try:
        with open(config_path, 'r') as f:
            content = f.read()
            print(f"配置文件内容预览:\n{content[:500]}...")

            if 'MobileViTv2Block' in content:
                print("⚠️ 警告：配置文件中包含 MobileViTv2Block 模块")
            else:
                print("✅ 配置文件中不包含 MobileViTv2Block 模块")
    except Exception as e:
        print(f"读取配置文件失败: {e}")

    # 加载模型
    model = YOLO(config_path)

    # 打印模型结构前几层
    print("\n模型结构摘要:")
    for name, module in model.model.named_children():
        print(f"模块: {name} - {type(module).__name__}")
        if hasattr(module, 'named_children'):
            for child_name, child_module in module.named_children():
                print(f"  └─ {child_name}: {type(child_module).__name__}")
        break  # 只打印第一个模块

    # 开始训练autodl-tmp/ljy/ultralytics/ultralytics/cfg/models/11/yolo11n.yaml
    results = model.train(
        data="/root/autodl-tmp/yolo11n_ljy/ultralytics/N-RDD2024/merged_dataset/data.yaml",  # 数据集配置文件的路径
        epochs=200,  # 训练的总轮数
        imgsz=640,  # 输入图像的大小
        batch=32,  # 批量大小
        device=0,  # 使用GPU设备（例如0,1,2,3或'cpu'）
        workers=8,  # 数据加载的工作进程数（减少避免网络问题）
        project="runs/detect",  # 保存训练结果的项目目录
        name="motor_bike_exp",  # 实验名称
        optimizer="SGD",  # 自动选择优化器
        lr0=0.01,  # 初始学习率
        patience=50,  # 早停耐心值
        save=True,  # 保存训练检查点和最终模型
        exist_ok=False, # 允许覆盖已有的实验目录
        cos_lr = True
    )


if __name__ == '__main__':
    main()